import math

class KNN:
    def __init__(self, k=3, task='classification'):
        self.k = k
        self.task = task
        self.X_train = None
        self.y_train = None

    def fit(self, X, y):
        self.X_train = X
        self.y_train = y
        return self

    def _euclidean_distance(self, a, b):
        distance = 0.0
        for i in range(len(a)):
            distance += (a[i] - b[i]) ** 2
        return math.sqrt(distance)

    def _get_k_neighbors(self, x):
        distances = []

        for i in range(len(self.X_train)):
            dist = self._euclidean_distance(x, self.X_train[i])
            distances.append((i, dist))

        for i in range(len(distances)):
            for j in range(i + 1, len(distances)):
                if distances[i][1] > distances[j][1]:
                    distances[i], distances[j] = distances[j], distances[i]

        return distances[:self.k]

    def _most_common(self, labels):
        count_dict = {}
        for label in labels:
            if label in count_dict:
                count_dict[label] += 1
            else:
                count_dict[label] = 1

        max_count = 0
        most_common_label = None
        for label, count in count_dict.items():
            if count > max_count:
                max_count = count
                most_common_label = label

        return most_common_label

    def predict(self, X):
        predictions = []

        for sample in X:
            neighbors = self._get_k_neighbors(sample)

            neighbor_labels = [self.y_train[idx] for idx, _ in neighbors]

            if self.task == 'classification':
                prediction = self._most_common(neighbor_labels)
                predictions.append(prediction)

            elif self.task == 'regression':
                total = 0.0
                for label in neighbor_labels:
                    total += label
                predictions.append(total / len(neighbor_labels))

        return predictions

    def predict_proba(self, X):
        if self.task != 'classification':
            raise ValueError("predict_proba 仅适用于分类任务")

        probabilities = []

        for sample in X:
            neighbors = self._get_k_neighbors(sample)
            neighbor_labels = [self.y_train[idx] for idx, _ in neighbors]

            label_count = {}
            for label in neighbor_labels:
                if label in label_count:
                    label_count[label] += 1
                else:
                    label_count[label] = 1

            total = len(neighbors)
            all_labels = sorted(set(self.y_train))
            prob_list = []
            for label in all_labels:
                count = label_count.get(label, 0)
                prob_list.append(count / total)
            probabilities.append(prob_list)

        return probabilities

    def score(self, X, y):
        predictions = self.predict(X)

        if self.task == 'classification':
            correct = 0
            for i in range(len(predictions)):
                if predictions[i] == y[i]:
                    correct += 1
            return correct / len(y)
        else:
            total_y = 0.0
            for value in y:
                total_y += value
            mean_y = total_y / len(y)

            ss_tot = 0.0
            for value in y:
                ss_tot += (value - mean_y) ** 2

            ss_res = 0.0
            for i in range(len(y)):
                ss_res += (y[i] - predictions[i]) ** 2

            if ss_tot == 0:
                return 0.0
            return 1 - (ss_res / ss_tot)